// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//  File: avtZoneDumpFilter.C
// ************************************************************************* //

#include <avtZoneDumpFilter.h>
#include <avtParallel.h>

#include <vtkCellData.h>
#include <vtkDataArray.h>
#include <vtkDataSet.h>
#include <vtkIntArray.h>
#include <vtkUnsignedIntArray.h>
#include <vtkUnsignedCharArray.h>

#ifdef PARALLEL
  #include <mpi.h>
#endif

#include <string>
#include <vector>

// ****************************************************************************
//  Method: avtZoneDumpFilter constructor
//
//  Programmer: harrison37 -- generated by xml2avt
//  Creation:   Tue Mar 13 09:00:34 PDT 2007
//
// ****************************************************************************

avtZoneDumpFilter::avtZoneDumpFilter()
{
}


// ****************************************************************************
//  Method: avtZoneDumpFilter destructor
//
//  Programmer: harrison37 -- generated by xml2avt
//  Creation:   Tue Mar 13 09:00:34 PDT 2007
//
//  Modifications:
//
// ****************************************************************************

avtZoneDumpFilter::~avtZoneDumpFilter()
{
}


// ****************************************************************************
//  Method:  avtZoneDumpFilter::Create
//
//  Programmer: harrison37 -- generated by xml2avt
//  Creation:   Tue Mar 13 09:00:34 PDT 2007
//
// ****************************************************************************

avtFilter *
avtZoneDumpFilter::Create()
{
    return new avtZoneDumpFilter();
}


// ****************************************************************************
//  Method:      avtZoneDumpFilter::SetAtts
//
//  Purpose:
//      Sets the state of the filter based on the attribute object.
//
//  Arguments:
//      a        The attributes to use.
//
//  Programmer: harrison37 -- generated by xml2avt
//  Creation:   Tue Mar 13 09:00:34 PDT 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::SetAtts(const AttributeGroup *a)
{
    atts = *(const ZoneDumpAttributes*)a;
        
    if(strcmp(atts.GetVariable().c_str(), "default") != 0)
    {
        SetActiveVariable(atts.GetVariable().c_str());
    }
}


// ****************************************************************************
//  Method: avtZoneDumpFilter::Equivalent
//
//  Purpose:
//      Returns true if creating a new avtZoneDumpFilter with the given
//      parameters would result in an equivalent avtZoneDumpFilter.
//
//  Programmer: harrison37 -- generated by xml2avt
//  Creation:   Tue Mar 13 09:00:34 PDT 2007
//
// ****************************************************************************

bool
avtZoneDumpFilter::Equivalent(const AttributeGroup *a)
{
    return (atts == *(ZoneDumpAttributes*)a);
}


// ****************************************************************************
//  Method: avtZoneDumpFilter::PreExecute
//
//  Purpose:  Clears zones array to prepare for execution.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::PreExecute()
{
    zones.clear();
}



// ****************************************************************************
//  Method: avtZoneDumpFilter::ExecuteData
//
//  Purpose:
//      Sends the specified input and output through the ZoneDump filter.
//
//  Arguments:
//      in_dr      The input data representation.
//
//  Returns:       The output data representation.
//
//  Programmer: harrison37 -- generated by xml2avt
//  Creation:   Tue Mar 13 09:00:34 PDT 2007
//
//  Modifications:
//    Cyrus Harrison, Mon Sep 24 11:30:58 PDT 2007
//    Pass correct zone id to GetOriginalLogicalIndices and ignore ghost 
//    zones
//
//    Kathleen Bonnell, Tue Jul 1 11:59:46 PDT 2008 
//    Removed unreferenced variables.
//
//    Eric Brugger, Tue Aug 19 09:53:43 PDT 2014
//    Modified the class to work with avtDataRepresentation.
//
// ****************************************************************************

avtDataRepresentation *
avtZoneDumpFilter::ExecuteData(avtDataRepresentation *in_dr)
{
    //
    // Get the VTK data set and domain number.
    //
    vtkDataSet *in_ds = in_dr->GetDataVTK();
    int domain = in_dr->GetDomain();

    if(atts.GetEnabled())
    {
        // Get the original cell number, the group_id, and value arrays.
        vtkDataArray *zone_ids = in_ds->GetCellData()->GetArray("avtOriginalCellNumbers");
        vtkDataArray *group_ids = in_ds->GetFieldData()->GetArray("group_id");
        vtkDataArray *var_vals = in_ds->GetCellData()->GetScalars();
        // get ghost zones array if it exists
        vtkUnsignedCharArray *gz_array = (vtkUnsignedCharArray *) in_ds
                                         ->GetCellData()->GetArray("avtGhostZones");
        unsigned char *gz_ptr = NULL;
        if(gz_array)
            gz_ptr = (unsigned char *)gz_array->GetPointer(0);
                                         
        // check for valid data
        if(!var_vals || !group_ids || !zone_ids)
        {
            EXCEPTION0(ImproperUseException);
        }

        // get the group id
        int group_id = (int) group_ids->GetTuple1(0);

        // get number of cells to process
        int ncells = in_ds->GetNumberOfCells();

        // find out the how the original zone #'s are stored
        bool  comps = zone_ids->GetNumberOfComponents() >1;
        int   zid;

        int ijk[3];

        // get requested dump bounds
        float vmin = atts.GetLowerBound();
        float vmax = atts.GetUpperBound();

        // loop through cells
        for(int i=0; i < ncells ;i++)
        {
            // dont process ghost zones
            if(gz_ptr && gz_ptr[i] == 1)
                continue;

            // get original zone #
            if(comps)
                zid = (int) zone_ids->GetComponent(i,1);
            else
                zid = (int) zone_ids->GetTuple1(i);

            // get value
            float val = var_vals->GetTuple1(i);

            // check if value falls within the desired range
            if(val >= vmin && val <= vmax )
            {
                // create a new zone info entry
                GetOriginalLogicalIndices(in_ds,zid,ijk);
                zones.push_back(ZoneInfo(group_id,domain,zid,
                                         ijk[0],ijk[1],ijk[2],
                                         val));
            }
        }
    }

    return in_dr;
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::PostExecute
//
//  Purpose:  Collects found zones at the root processor and writes zone info
//            to an output file. (Serial case simply writes to output file)
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::PostExecute()
{
    // skip dump if not enabled
    if(!atts.GetEnabled())
        return;

#ifdef PARALLEL
    // loop index
    int i;

    // get the number of processors and the current processor id
    int nprocs = PAR_Size();
    int procid = PAR_Rank();

    // get the number of zone infos to send
    int  n_snd_zones = zones.size();
    // size of each zone info
    int  zinfo_size = ZoneInfo::PackedSize();
    // calculate the total send message size
    int  snd_msg_size = n_snd_zones * zinfo_size;
    // send buffer
    unsigned char *snd_msg = NULL;

    // vars for the root processor
    // size of the gather message
    int  rcv_msg_size = 0;
    // holds size of the msg from each other processor
    int *rcv_count = NULL;
    // holds the displacement of the msg from each other processor
    int *rcv_disp  = NULL;
    // receive buffer
    unsigned char *rcv_msg = NULL;

    if(procid == 0)
    {
        // allocate space for these for root proc only
        rcv_count = new int[nprocs];
        rcv_disp  = new int[nprocs];
    }

    // gather message sizes from all procs to root proc
    MPI_Gather(&snd_msg_size,1, MPI_INT,
               rcv_count, 1, MPI_INT,
               0, VISIT_MPI_COMM);

    // find message offsets and total rcv size
    if(procid == 0)
    {
        rcv_disp[0]  = 0;
        rcv_msg_size = rcv_count[0];

        for( i=1; i<nprocs;i++)
        {
            rcv_disp[i]   = rcv_count[i-1]  +  rcv_disp[i-1];
            rcv_msg_size += rcv_count[i];
        }
    }

    // get total # of zone infos
    int nrcv_zones = rcv_msg_size / zinfo_size;

    // create msg to send to the root proc
    if(snd_msg_size > 0)
    {
        snd_msg= new unsigned char[snd_msg_size];
        unsigned char *snd_msg_ptr = snd_msg;

        // pack zone infos
        for(i=0; i < n_snd_zones; i++)
        {
            zones[i].Pack(snd_msg_ptr);
            snd_msg_ptr+= zinfo_size;
        }
    }

    if(procid == 0 && rcv_msg_size > 0)
    {
        // create the rcv buffer for the root proc
        rcv_msg = new unsigned char[rcv_msg_size];
    }

    // gather all zone infos
    MPI_Gatherv(snd_msg, snd_msg_size, MPI_UNSIGNED_CHAR,
                rcv_msg, rcv_count, rcv_disp, MPI_UNSIGNED_CHAR,
                0,VISIT_MPI_COMM);

    if(procid == 0 )
    {
        // unpack all rcvd zones
        std::vector<ZoneInfo> rcv_zones(nrcv_zones);

        unsigned char *rcv_msg_ptr = rcv_msg;

        for( i = 0; i < nrcv_zones; i++)
        {
            rcv_zones[i].Unpack(rcv_msg_ptr);
            rcv_msg_ptr += zinfo_size;
        }

        // save all zones
        SaveOutput(atts.GetOutputFile(),rcv_zones);
    }

    // cleanup

    if(snd_msg)
        delete[] snd_msg;

    if(rcv_msg)
        delete[] rcv_msg;

    if(rcv_count)
        delete[] rcv_count;

    if(rcv_disp)
        delete[] rcv_disp;
#else

    // for serial case, simply dump out zones found during the exe pass.
    SaveOutput(atts.GetOutputFile(),zones);

#endif
}



// ****************************************************************************
//  Method: avtZoneDumpFilter::ModifyContract
//
//  Purpose:  Requests Original Zone Numbers
//
//  Programmer: Cyrus Harrison
//  Creation:   March 13, 2007
//
// ****************************************************************************

avtContract_p
avtZoneDumpFilter::ModifyContract(avtContract_p contract)
{
    avtContract_p rv = contract;
    // get original zone numbers
    rv->GetDataRequest()->TurnZoneNumbersOn();
    // get original zone indices (for finding logical coords)
    rv->GetDataRequest()->SetNeedStructuredIndices(true);
    return rv;
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::GetOriginalLogicalIndices
//
//  Purpose:  Gets original logical indices
//
//  Notes:    Adapted from vtkVisItUtility::GetLogicalIndices
//
//  Programmer: Cyrus Harrison
//  Creation:   March 14, 2007
//
//  Modifications:
//    Cyrus Harrison, Mon Apr 30 10:50:12 PDT 2007
//    Fixed spelling mistake. 
//
//    Cyrus Harrison, Mon Sep 24 07:59:55 PDT 2007
//    Fixed error with j index for the 2d case.
//
// ****************************************************************************

void 
avtZoneDumpFilter::GetOriginalLogicalIndices(vtkDataSet *ds,
                                            int zone,
                                            int ijk[3])
{
    int dims[3] = {0,0,0};
    int base[3] = {0,0,0};

    // get dims
    vtkUnsignedIntArray *orig_dims = (vtkUnsignedIntArray*)ds->GetFieldData()
                                 ->GetArray("avtOriginalStructuredDimensions");

    // if original dims array exists, use to set dims
    if(orig_dims)
    {
        dims[0] = orig_dims->GetValue(0);
        dims[1] = orig_dims->GetValue(1);
        dims[2] = orig_dims->GetValue(2);

        dims[0] = (dims[0]-1 > 0 ? dims[0]-1 : 1);
        dims[1] = (dims[1]-1 > 0 ? dims[1]-1 : 1);
    }

    // get base offset
    vtkIntArray *base_offset = (vtkIntArray*)ds->GetFieldData()
                                                    ->GetArray("base_index");
    if(base_offset)
    {
        base[0] = base_offset->GetValue(0);
        base[1] = base_offset->GetValue(1);
        base[2] = base_offset->GetValue(2);
    }


    // find the logical coords
    if (dims[2] == 1)
    {
        ijk[0] = (zone % dims[0]) + base[0];
        ijk[1] = ((zone / dims[0]) % dims[1]) + base[1];
        ijk[2] = 0;
    }
    else 
    {
        ijk[0] = (zone % dims[0])          + base[0];
        ijk[1] = ((zone / dims[0]) % dims[1]) + base[1];
        ijk[2] = (zone / (dims[0] * dims[1])) + base[2];
    }

    ijk[0] = ijk[0] < 0 ? 0 : ijk[0];
    ijk[1] = ijk[1] < 0 ? 0 : ijk[1];
    ijk[2] = ijk[2] < 0 ? 0 : ijk[2];
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::SaveOutput
//
//  Purpose:  Saves zone info to a file.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::SaveOutput(const std::string &fname, std::vector<ZoneInfo> &zones)
{
    // loop index
    int i;

    // get # of found zones
    int nzones = (int)zones.size();

    // open output file
    ofstream outs(fname.c_str());
    if(outs.fail())
        EXCEPTION0(ImproperUseException);

    // write general header info
    outs << "# VisIt Zone Dump  " << endl;
    outs << "# Number of Zones: " << nzones << endl;
    outs << "# Fields:" <<endl;
    outs << "block\tdomain\tzone\ti\tj\tk\t"
         << atts.GetVariable() << endl;

    // dump each found zone
    for(i = 0; i < nzones; i++)
        zones[i].Print(outs,"\t");

    // close the output file
    outs.close();
}


// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo default constructor
//
//  Purpose:
//      Initializes a Zone Info object.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

avtZoneDumpFilter::ZoneInfo::ZoneInfo()
{
    this->groupId  = 0;
    this->domainId = 0;
    this->zone     = 0;
    this->i        = 0;
    this->j        = 0;
    this->k        = 0;
    this->value    = 0;
}


// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo constructor
//
//  Purpose:
//      Initializes a Zone Info object.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

avtZoneDumpFilter::ZoneInfo::ZoneInfo
(int g, int d, int z, int i, int j, int k, float val)
{
    this->groupId  = g;
    this->domainId = d;
    this->zone     = z;
    this->i        = i;
    this->j        = j;
    this->k        = k;
    this->value    = val;
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo destructor
//
//  Purpose:
//      Cleans up a Zone Info object.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

avtZoneDumpFilter::ZoneInfo::~ZoneInfo()
{
    ;
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo::Print()
//
//  Purpose:
//      Prints zone info to an output stream.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::ZoneInfo::Print(ostream &os, const std::string &sep) const
{
    os << groupId  << sep
       << domainId << sep
       << zone     << sep
       << i        << sep
       << j        << sep
       << k        << sep
       << value    << endl;
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo::Pack()
//
//  Purpose:
//      Packs zone info into a buffer.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::ZoneInfo::Pack(unsigned char *buff)
{
    // pack int fields
    int *buff_int_ptr = (int*)buff;

    buff_int_ptr[0] = groupId;
    buff_int_ptr[1] = domainId;
    buff_int_ptr[2] = zone;
    buff_int_ptr[3] = i;
    buff_int_ptr[4] = j;
    buff_int_ptr[5] = k;

    // pack float field
    float *buff_flt_ptr = (float*)(buff_int_ptr + 6);

    buff_flt_ptr[0] = value;
}


// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo::Unpack()
//
//  Purpose:
//      Unpacks zone info from a buffer.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 2, 2007
//
// ****************************************************************************

void
avtZoneDumpFilter::ZoneInfo::Unpack(unsigned char *buff)
{
    // unpack int fields
    int *buff_int_ptr = (int*)buff;

    groupId   = buff_int_ptr[0];
    domainId  = buff_int_ptr[1];
    zone      = buff_int_ptr[2];
    i         = buff_int_ptr[3];
    j         = buff_int_ptr[4];
    k         = buff_int_ptr[5];

    // unpack float field
    float *buff_flt_ptr = (float*)(buff_int_ptr + 6);

    value = buff_flt_ptr[0];
}

// ****************************************************************************
//  Method: avtZoneDumpFilter::ZoneInfo::operator ==
//
//  Purpose:
//      Compares two zone info objects.
//
//  Programmer: Cyrus Harrison
//  Creation:   April 3, 2007
//
// ****************************************************************************

bool
avtZoneDumpFilter::ZoneInfo::operator==(const ZoneInfo &b)
{
    if( this->groupId  != b.groupId) return false;
    if( this->domainId != b.domainId ) return false;
    if( this->zone     != b.zone) return false;
    if( this->i        != b.i) return false;
    if( this->j        != b.j) return false;
    if( this->k        != b.k) return false;
    if( this->value    != b.value) return false;

    return true;

}


